#!/usr/bin/python3.7
# -*- coding: utf-8 -*-
# create by 火苗999℃
import random
from draw_maze import * 
from maze_def import *
from maze import *

# Randomized Kruskal's algorithm
# This algorithm is a randomized version of Kruskal's algorithm.
# 1.Create a list of all walls, and create a set for each cell, each containing just that one cell.
# 2.For each wall, in some random order:
#    1.If the cells divided by this wall belong to distinct sets:
#        1.Remove the current wall.
#        2.Join the sets of the formerly divided cells.

# 随机的Kruskal算法
# 这个算法是Kruskal算法的随机化版本。
# 1。创建所有墙壁的列表，并为每个单元格创建一个集合，每个集合只包含一个单元格。
# 2。对于每一面墙，以一些随机的顺序:
#   1。如果由这个壁分隔的细胞属于不同的集合:
#       1。移除当前的墙。
#       2。加入以前分裂的细胞组。


# 算法
def kruskal_maze(levels, rows, cols):
    if 0 == levels % 2:
        levels+=1
    if 0 == rows % 2:
        rows+=1
    if 0 == cols % 2:
        cols+=1
    # 初始化全为墙
    # maze_map=[[[ WALL for z in range(levels)]for y in range(rows)] for x in range(cols)]
    maze_map=maze_init_draw(levels,rows,cols)
    # 设置起点
    x=1
    y=1
    z=0
    # 格子列表
    gridlist=[]
    gridlist.append((x,y,z))
    # 单元格集合
    collection =[]
    # 墙壁的列表
    wallList=[]
    # 标记格子未访问
    for zi in range(0, levels, 1):
        for xi in range(0, cols, 1):
            for yi in range(0, rows, 1):
                if xi%2 == 1 and yi%2==1 and zi%2==0:
                    # 标记为格子
                    maze_map[xi][yi][zi]=CELL_NO_VISIT
                    collection.append([(xi,yi,zi)])
                    continue
                # 跳过不能打通的墙
                if xi == 0:
                    continue
                if yi == 0:
                    continue
                if xi == cols-1:
                    continue
                if yi == rows-1:
                    continue
                if zi%2==0 and xi%2==0 and yi%2==0:
                    continue
                if zi%2==1 and not (xi%2==1 and yi%2==1):
                    continue
                # 可打通的墙壁
                if zi%2==1:
                    wallList.append((xi,yi,zi, UD))
                    continue
                if xi%2==0:
                    wallList.append((xi,yi,zi, LR))
                    continue
                if yi%2==0:
                    wallList.append((xi,yi,zi, FB))
                    continue
                print("error")

    while wallList:
        # 随机选一个墙
        x,y,z,d = random.choice(wallList)
        # 每个墙随机到一次
        wallList.remove((x,y,z,d))
        # a,b相邻的两个格子
        if d == UD: # 竖墙
            grida = (x,y,z-1)
            gridb = (x,y,z+1)
        elif d==LR :  # 横墙
            grida = (x-1,y,z)
            gridb = (x+1,y,z)
        elif d==FB:
            grida = (x,y-1,z)
            gridb = (x,y+1,z)

        colla = []
        collb = []
        for coll in collection:
            if grida in coll:
                colla = coll
            if gridb in coll:
                collb = coll
        # 设置访问过
        maze_map[grida[0]][grida[1]][grida[2]] = CELL_VISIT
        maze_map[gridb[0]][gridb[1]][gridb[2]] = CELL_VISIT
        # 
        if colla != collb:
            # 打通墙
            maze_map[x][y][z] = NOWALL
            # 合并集合
            coll = colla+collb
            collection.remove(colla)
            collection.remove(collb)
            collection.append(coll)
    return maze_map


# 随机格子
def kruskal_maze_demo(levels, rows, cols):
    if 0 == levels % 2:
        levels+=1
    if 0 == rows % 2:
        rows+=1
    if 0 == cols % 2:
        cols+=1
    # 初始化全为墙
    maze_map = maze_init_draw(levels, rows, cols)
    # 单元格集合
    collection =[]
    # 墙壁的列表
    wallList=[]
    # 标记格子和墙
    for zi in range(0, levels):
        for xi in range(0, cols):
            for yi in range(0, rows):
                if maze_map[xi][yi][zi] == CELL_NO_VISIT:
                    # 标记为格子
                    collection.append([(xi,yi,zi)])
                    continue
                # 跳过不能打通的墙
                if maze_map[xi][yi][zi] == STEEL:
                    continue
                # 可打通的墙壁
                if zi%2==1:
                    wallList.append((xi,yi,zi, UD))
                    continue
                if xi%2==0 and yi%2==1:
                    wallList.append((xi,yi,zi, LR))
                    continue
                if xi%2==1 and yi%2==0:
                    wallList.append((xi,yi,zi, FB))
                    continue
                print("error")

    while True:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                return
        if  wallList:
            # 随机选一个墙
            nowx,nowy,nowz,d = random.choice(wallList)
            # 每个墙随机到一次
            wallList.remove((nowx,nowy,nowz,d))
            # a,b相邻的两个格子
            if d == UD: # 分隔上下两个格子的墙
                grida = (nowx,nowy,nowz-1)
                gridb = (nowx,nowy,nowz+1)
            elif d==LR :  # 分隔左右两个格子的墙
                grida = (nowx-1,nowy,nowz)
                gridb = (nowx+1,nowy,nowz)
            elif d==FB: # 分隔前后两个格子的墙
                grida = (nowx,nowy-1,nowz)
                gridb = (nowx,nowy+1,nowz)
            # 查找两个墙所属的集合
            colla = []
            collb = []
            for coll in collection:
                if grida in coll:
                    colla = coll
                if gridb in coll:
                    collb = coll
            maze_map[nowx][nowy][nowz] = WALL_VISIT # 标记墙壁已经访问
            # 两个集合不联通
            if colla != collb:
                # 打通墙
                maze_map[nowx][nowy][nowz] = NOWALL
                # a,b相邻的两个格子
                if d == UD: # 分隔上下两个格子的墙
                    grida = (nowx,nowy,nowz-1)
                    gridb = (nowx,nowy,nowz+1)
                    # 设置两个格子为访问过
                    if maze_map[grida[0]][grida[1]][grida[2]] == STAIRS_D:
                        maze_map[grida[0]][grida[1]][grida[2]] = STAIRS_UD
                    else:
                        maze_map[grida[0]][grida[1]][grida[2]] = STAIRS_U
                    
                    if maze_map[gridb[0]][gridb[1]][gridb[2]] == STAIRS_U:
                        maze_map[gridb[0]][gridb[1]][gridb[2]] = STAIRS_UD
                    else:
                        maze_map[gridb[0]][gridb[1]][gridb[2]] = STAIRS_D
                elif d==LR :  # 分隔左右两个格子的墙
                    grida = (nowx-1,nowy,nowz)
                    gridb = (nowx+1,nowy,nowz)
                    # 设置两个格子为访问过
                    if maze_map[grida[0]][grida[1]][grida[2]] == CELL_NO_VISIT:
                        maze_map[grida[0]][grida[1]][grida[2]] = CELL_VISIT
                    if maze_map[gridb[0]][gridb[1]][gridb[2]] == CELL_NO_VISIT:
                        maze_map[gridb[0]][gridb[1]][gridb[2]] = CELL_VISIT
                elif d==FB: # 分隔前后两个格子的墙
                    grida = (nowx,nowy-1,nowz)
                    gridb = (nowx,nowy+1,nowz)
                    # 设置两个格子为访问过
                    if maze_map[grida[0]][grida[1]][grida[2]] == CELL_NO_VISIT:
                        maze_map[grida[0]][grida[1]][grida[2]] = CELL_VISIT
                    if maze_map[gridb[0]][gridb[1]][gridb[2]] == CELL_NO_VISIT:
                        maze_map[gridb[0]][gridb[1]][gridb[2]] = CELL_VISIT
                # 合并集合
                coll = colla+collb
                collection.remove(colla)
                collection.remove(collb)
                collection.append(coll)
        # 
        draw_maze(maze_map, levels, rows, cols, (nowx,nowy,nowz), not wallList)
        time_passed = clock.tick(5)

        pygame.display.update()
    return 



# main
if __name__ == "__main__":
    '''main'''
    kruskal_maze_demo(5, 11, 15)
